import base64
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding


def pad_key_iv(key, iv):
    key = key.ljust(32, b'\0')[:32]
    iv = iv.ljust(16, b'\0')[:16]
    return key, iv


def encrypt(plaintext, key, iv):
    # 将字符串明文转换为字节串
    plaintext_bytes = plaintext.encode('utf-8')

    key, iv = pad_key_iv(key, iv)

    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    encryptor = cipher.encryptor()

    padder = padding.PKCS7(algorithms.AES.block_size).padder()
    padded_data = padder.update(plaintext_bytes) + padder.finalize()

    ciphertext = encryptor.update(padded_data) + encryptor.finalize()

    return base64.b64encode(ciphertext).decode('utf-8')


def decrypt(ciphertext, key, iv):
    key, iv = pad_key_iv(key, iv)

    ciphertext_bytes = base64.b64decode(ciphertext)

    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
    decryptor = cipher.decryptor()

    padded_plaintext = decryptor.update(ciphertext_bytes) + decryptor.finalize()

    unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
    plaintext_bytes = unpadder.update(padded_plaintext) + unpadder.finalize()

    # 将字节串转换回字符串
    return plaintext_bytes.decode('utf-8')


# 示例用法
key = b'shortkey'
iv = b'shortiv'
plaintext = 'This is a secret message.'

ciphertext = encrypt(plaintext, key, iv)
print("Ciphertext (Base64):", ciphertext)

decrypted_text = decrypt(ciphertext, key, iv)
print("Decrypted text:", decrypted_text)
